package offer;
import java.util.Arrays;
import java.util.PriorityQueue;
import java.util.Queue;

/**
 * 最小的K个数 - Offer40
 */
public class LeastNumbersK {
    public static void main(String[] args) {
        int[] arr = new int[] { 0, 1, 2, 1 };
        int[] res = getLeastNumbers(arr, 1);
        System.out.println(Arrays.toString(res));
    }

    // 注意不要使用小根堆, 这样做会把全部元素入堆
    // 维持一个大小为k的窗口, 每次把大元素出堆(和大根堆堆顶比较), 最后留下k个元素即答案
    public static int[] getLeastNumbers(int[] arr, int k) {
        int[] res = new int[k];
        if (k == 0) {
            return res;
        }
        // 转为大顶堆的优先队列
        Queue<Integer> queue = new PriorityQueue<>((i1, i2) -> i2 - i1);
        // 初始化 填满k个
        for (int i = 0; i < k; i++) {
            queue.offer(arr[i]);
        }
        for (int i = k; i < arr.length; i++) {
            if (queue.peek() > arr[i]) {
                queue.poll(); // 弹出
                queue.offer(arr[i]); // 压入小值
            }
        }
        for (int i = 0; i < res.length; i++) {
            res[i] = queue.poll();
        }
        return res;
    }
}
